import argparse
import sys
from typing import Callable, Dict, Any, Optional, List
from dataclasses import dataclass, field

from engine_server.config.base import BaseConfig, ServerConfig
from vllm.utils import FlexibleArgumentParser
from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
from engine_server.utils.logger import run_log


def _add_argument_to_list(arg_list: list, key: str, value: Any):
    if isinstance(value, bool):
        if value:
            arg_list.append(f"--{key}")
    elif isinstance(value, list):
        if value:
            arg_list.append(f"--{key}")
            for item in value:
                arg_list.append(str(item))
    else:
        arg_list.append(f"--{key}")
        arg_list.append(str(value))


def default_checker(value: Any) -> bool:
    return True


def dp_rank_checker(value: Any) -> bool:
    return value > 0


@dataclass
class VLLMConfig(BaseConfig):
    args: Optional[argparse.Namespace] = None
    mapping: Dict[str, str] = field(default_factory=lambda: {
        'dp_rank': 'data_parallel_rank',
    })
    checkers: Dict[str, Callable[[Any], bool]] = field(default_factory=lambda: {
        'dp_rank': dp_rank_checker,
    })
    
    def validate(self):
        super().validate()
        if self.args is not None:
            validate_parsed_serve_args(self.args)

    def convert(self):
        super().convert()
        arg_list = self._get_param_list()
        run_log.info(f'engine server parsed arg_list: {arg_list}')

        sys.argv = ["serve"] + arg_list
        
        parser = FlexibleArgumentParser(description="vLLM parser")
        parser = make_arg_parser(parser)
        self.args = parser.parse_args()

    def get_args(self) -> argparse.Namespace:
        return self.args

    def get_server_config(self) -> ServerConfig:
        return self.server_config
    
    def _should_include_param(self, key: str, value: Any) -> bool:
        checker = self.checkers.get(key, default_checker)
        return checker(value)

    def _get_param_list(self) -> List[str]:
        processed_args = []

        for server_key, vllm_key in self.mapping.items():
            if hasattr(self.server_config, server_key):
                value = getattr(self.server_config, server_key)
                if self._should_include_param(server_key, value):
                    formatted_key = vllm_key.replace('_', '-')
                    _add_argument_to_list(processed_args, formatted_key, value)

        for key, value in self.kv_config.items():
            if self._should_include_param(key, value):
                formatted_key = key.replace('_', '-')
                _add_argument_to_list(processed_args, formatted_key, value)
        
        return processed_args
        


